# coding: utf-8
import datetime

import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F
from rich.console import Console
from rich.progress import track
from torchvision import datasets, models, transforms

console = Console()


def print(*args, **kwargs):
    console.print(*args, **kwargs)


def now():
    return datetime.datetime.now().strftime("%Y-%m-%d %H:%M")


# Download the CIFAR10 dataset
data_set = datasets.CIFAR10(
    "data/", download=True, train=True, transform=transforms.ToTensor()
)

# Split the dataset into train and test sets
train_set, test_set = torch.utils.data.random_split(data_set, [40000, 10000])

# Create data loaders
train_loader = torch.utils.data.DataLoader(train_set, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=64, shuffle=True)


# # Create a model
# model = models.resnet50(weights=models.resnet.ResNet50_Weights.DEFAULT)

# # Visualize the model


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 64)
        self.fc2 = nn.Linear(64, 32)
        self.fc3 = nn.Linear(32, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))  # 1x32x32 -> 6x28x28
        x = self.pool(x)  # 6x28x28 -> 6x14x14
        x = F.relu(self.conv2(x))  # 6x14x14 -> 16x10x10
        x = self.pool(x)  # 16x10x10 -> 16x5x5
        x = x.view(-1, 16 * 5 * 5)  # 16x5x5 -> 400
        x = F.relu(self.fc1(x))  # 400 -> 64
        x = F.relu(self.fc2(x))  # 64 -> 32
        x = self.fc3(x)  # 32 -> 10
        return x

    # Define the training loop
    def fit(
        self,
        train_loader,
        loss_fn=nn.CrossEntropyLoss(),
        optimizer=torch.optim.Adam,
        epochs=10,
        lr=0.001,
    ):
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.to(device)
        optimizer = optimizer(self.parameters(), lr=lr)
        epoch_loss = []
        epoch_train_acc = []
        epoch_test_acc = []
        for epoch in range(epochs):
            self.train()
            running_loss = 0.0
            for inputs, labels in track(
                train_loader, description=f"Epoch {epoch+1}/{epochs}"
            ):
                optimizer.zero_grad()
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = self(inputs)
                loss = loss_fn(outputs, labels)
                loss.backward()
                optimizer.step()
                running_loss += loss.item()

            train_loss = running_loss / len(train_loader)
            train_acc = self.accuracy(train_loader)
            test_acc = self.accuracy(test_loader)
            epoch_loss.append(train_loss)
            epoch_train_acc.append(train_acc)
            epoch_test_acc.append(test_acc)
            plt.clf()
            plt.subplot(1, 2, 1)
            plt.title(f"Loss: {train_loss:.3f}")
            plt.plot(epoch_loss, label="train_loss", marker="o")
            plt.xlabel("Epoch")
            plt.ylabel("Loss")
            plt.xticks(range(1, epochs + 1))
            plt.legend()
            plt.subplot(1, 2, 2)
            plt.title(f"Accuracy: {train_acc:.3f} / {test_acc:.3f}")
            plt.plot(epoch_train_acc, label="train_acc", marker="o")
            plt.plot(epoch_test_acc, label="test_acc", marker="o")
            plt.xlabel("Epoch")
            plt.ylabel("Accuracy")
            plt.xticks(range(1, epochs + 1))
            plt.legend()

        plt.tight_layout()
        plt.savefig(f"001.intro.{now()}.png")
        plt.show()

    def accuracy(self, loader):
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.to(device)
        self.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for inputs, labels in loader:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = self(inputs)
                _, predicted = torch.max(outputs, dim=1)
                total += labels.shape[0]
                correct += int((predicted == labels).sum())
        return correct / total


model = Net()
model.fit(train_loader, epochs=20, lr=0.001)
